import os
import random
import argparse
import librosa
import soundfile as sf

def check_sample_rate(directory,num_samples=20,extensions=['.wav','.flac','.mp3']):
    """
    随机检查指定目录中音频文件的采样率
    
    参数:
        directory (str): 要检查的目录路径
        num_samples (int): 随机检查的文件数量
        extensions (list): 要检查的音频文件扩展名列表
    """

        # 收集所有音频文件
    audio_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if any(file.lower().endswith(ext) for ext in extensions):
                audio_files.append(os.path.join(root, file))
    
    if not audio_files:
        print(f"在目录 {directory} 中未找到音频文件")
        return
    
    # 随机选择文件
    if len(audio_files) > num_samples:
        selected_files = random.sample(audio_files, num_samples)
    else:
        selected_files = audio_files
        print(f"目录中只有 {len(audio_files)} 个音频文件，将检查所有文件")
    
    print(f"开始检查 {len(selected_files)} 个随机音频文件的采样率...\n")
    
    sample_rates = {}
    problem_files = []
    
    for i, file_path in enumerate(selected_files, 1):
        try:
            # 使用librosa获取采样率 (不会加载整个音频文件)
            sr = librosa.get_samplerate(file_path)

            
            print(f"{i}/{len(selected_files)}: {os.path.basename(file_path)} - 采样率: {sr} Hz")
            
            # 统计不同采样率的出现次数
            if sr not in sample_rates:
                sample_rates[sr] = 0
            sample_rates[sr] += 1
            
        except Exception as e:
            print(f"{i}/{len(selected_files)}: {os.path.basename(file_path)} - 错误: {str(e)}")
            problem_files.append(file_path)
    
    # 打印统计结果
    print("\n" + "="*50)
    print("采样率统计:")
    for sr, count in sample_rates.items():
        print(f"  {sr} Hz: {count} 个文件")
    
    if problem_files:
        print(f"\n遇到问题的文件 ({len(problem_files)} 个):")
        for file in problem_files:
            print(f"  {file}")
    
    # 检查是否所有文件都是48kHz
    if 48000 in sample_rates and sample_rates[48000] == len(selected_files) - len(problem_files):
        print("\n✓ 所有检查的文件都是48kHz采样率")
    else:
        print("\n⚠ 注意: 不是所有文件都是48kHz采样率")
        if 48000 in sample_rates:
            print(f"  只有 {sample_rates[48000]}/{len(selected_files)} 个文件是48kHz")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='随机检查音频文件的采样率')
    parser.add_argument('directory', help='要检查的目录路径')
    parser.add_argument('-n', '--num-samples', type=int, default=20, 
                       help='随机检查的文件数量 (默认: 20)')
    
    args = parser.parse_args()
    
    if not os.path.isdir(args.directory):
        print(f"错误: 目录 '{args.directory}' 不存在")
        exit(1)
    
    check_sample_rate(args.directory, args.num_samples)